class WS2812Screen:
    def __init__(self, led, bufferWidth, bufferHeight):
        self.led = led
        self.index = 0
        self.w = bufferWidth
        self.h = bufferHeight

    def draw_pixel(self, x, y, rgb, bright=0.5):
        if bright > 1:
            bright = 1
        elif bright < 0:
            bright = 0

        if 0 <= x < self.w and 0 <= y < self.h:  # 判断xy是否超出屏幕范围
            if 0 <= x <= 7:  # 因为灯板的安装方式，导致led索引排布不规律，需要单独判断
                self.index = (199 - x) + (7 - y) * 8
            elif 8 <= x <= 15:
                self.index = (143 - x) + (7 - y) * 8
            elif 16 <= x <= 23:
                self.index = (87 - x) + (7 - y) * 8
            elif 24 <= x <= 31:
                self.index = (31 - x) + (7 - y) * 8
        if -8 <= x < 0:  # 屏幕左边超出屏幕范围，为了方便做动画出入
            self.index = (255 - x) + (7 - y) * 8
        if self.w <= x <= 39:  # 屏幕右边超出屏幕范围，为了方便做动画出入
            self.index = (-25 - x) + (7 - y) * 8
        self.led[self.index] = (int(rgb[0] * bright), int(rgb[1] * bright), int(rgb[2] * bright))  # 给RGB乘上一个系数，控制亮度

    def draw_num(self, x, y, num, rgb, bright=0.5):
        if num == 0:
            for h in range(5):
                self.draw_pixel(x, y + h, rgb, bright)
                self.draw_pixel(x + 2, y + h, rgb, bright)
            self.draw_pixel(x + 1, y, rgb, bright)
            self.draw_pixel(x + 1, y + 4, rgb, bright)

        elif num == 1:
            for h in range(5):
                self.draw_pixel(x + 1, y + h, rgb, bright)
            self.draw_pixel(x, y + 1, rgb, bright)
            self.draw_pixel(x, y + 4, rgb, bright)
            self.draw_pixel(x + 2, y + 4, rgb, bright)

        elif num == 2:
            for w in range(3):
                self.draw_pixel(x + w, y, rgb, bright)
                self.draw_pixel(x + w, y + 2, rgb, bright)
                self.draw_pixel(x + w, y + 4, rgb, bright)
            self.draw_pixel(x + 2, y + 1, rgb, bright)
            self.draw_pixel(x, y + 3, rgb, bright)

        elif num == 3:
            for w in range(3):
                self.draw_pixel(x + w, y, rgb, bright)
                self.draw_pixel(x + w, y + 2, rgb, bright)
                self.draw_pixel(x + w, y + 4, rgb, bright)
            self.draw_pixel(x + 2, y + 1, rgb, bright)
            self.draw_pixel(x + 2, y + 3, rgb, bright)

        elif num == 4:
            for h in range(5):
                self.draw_pixel(x + 2, y + h, rgb, bright)
            for h1 in range(3):
                self.draw_pixel(x, y + h1, rgb, bright)
            self.draw_pixel(x + 1, y + 2, rgb, bright)

        elif num == 5:
            for w in range(3):
                self.draw_pixel(x + w, y, rgb, bright)
                self.draw_pixel(x + w, y + 2, rgb, bright)
                self.draw_pixel(x + w, y + 4, rgb, bright)
            self.draw_pixel(x, y + 1, rgb, bright)
            self.draw_pixel(x + 2, y + 3, rgb, bright)

        elif num == 6:
            for w in range(3):
                self.draw_pixel(x + w, y, rgb, bright)
                self.draw_pixel(x + w, y + 2, rgb, bright)
                self.draw_pixel(x + w, y + 4, rgb, bright)
            self.draw_pixel(x, y + 1, rgb, bright)
            self.draw_pixel(x, y + 3, rgb, bright)
            self.draw_pixel(x + 2, y + 3, rgb, bright)

        elif num == 7:
            for h in range(5):
                self.draw_pixel(x + 2, y + h, rgb, bright)
            self.draw_pixel(x, y, rgb, bright)
            self.draw_pixel(x + 1, y, rgb, bright)

        elif num == 8:
            for h in range(5):
                self.draw_pixel(x, y + h, rgb, bright)
                self.draw_pixel(x + 2, y + h, rgb, bright)
            self.draw_pixel(x + 1, y, rgb, bright)
            self.draw_pixel(x + 1, y + 2, rgb, bright)
            self.draw_pixel(x + 1, y + 4, rgb, bright)

        elif num == 9:
            for w in range(3):
                self.draw_pixel(x + w, y, rgb, bright)
                self.draw_pixel(x + w, y + 2, rgb, bright)
                self.draw_pixel(x + w, y + 4, rgb, bright)
            self.draw_pixel(x, y + 1, rgb, bright)
            self.draw_pixel(x + 2, y + 1, rgb, bright)
            self.draw_pixel(x + 2, y + 3, rgb, bright)

    def draw_img(self, x, y, imgData, bright=0.5):
        for i in range(len(imgData)):
            for j in range(len(imgData[0])):
                if i >= self.h or j >= self.w:  # 超出屏幕范围的图像数据，不绘制(跳过)
                    continue
                self.draw_pixel(x + j, y + i, imgData[i][j], bright)

    def fill_screen(self, rgb=(0, 0, 0)):
        self.led.fill(rgb)

    def show(self):
        self.led.write()

    def drawNumber(self, number, x, y, width, rgb, bright):
         char1 = str(number)
         for item in [c for c in char1]:
            if (item == '.'):
                self.draw_pixel(x, y+4, rgb, bright)
                x = x + 2
            else:
                self.draw_num(x, y, int(item), rgb, bright)
                x = x + width
          
    def drawNumberAligned(self, number, rgb, bright, aligned="left"):
         char1 = str(number)
         length = len(char1)
         dotCount = char1.count('.')
         totalWitdh= (length-dotCount)*4 + dotCount*2
         if (aligned == "left"):
             self.drawNumber(number, 1, 1, 4, rgb, bright)
         if (aligned  == "right"):
             self.drawNumber(number, self.w - totalWitdh , 1, 4, rgb, bright)
         if (aligned  == "center"):
             self.drawNumber(number, int ((self.w - totalWitdh) / 2 ), 1, 4, rgb, bright)
          
           


